#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
多地域大规模GPU服务器部署脚本
支持5个地域并行部署，最大化部署能力
"""

import requests
import json
import time
import sys
import argparse
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Dict, Any
import asyncio

class MultiRegionDeploymentManager:
    def __init__(self, base_url: str = "http://localhost:8080"):
        self.base_url = base_url
        
        # 支持的地域配置
        self.regions = {
            "ap-beijing": {
                "name": "北京",
                "app_id": "app-oefac6i2",
                "api_limit_per_second": 20
            },
            "ap-nanjing": {
                "name": "南京", 
                "app_id": "app-9qle417n",
                "api_limit_per_second": 20
            },
            "ap-guangzhou": {
                "name": "广州",
                "app_id": "app-9diy38hx", 
                "api_limit_per_second": 20
            },
            "ap-shanghai": {
                "name": "上海",
                "app_id": "app-d8orpkas",
                "api_limit_per_second": 20
            },
            "ap-chongqing": {
                "name": "重庆",
                "app_id": "app-h5rw7hfh",
                "api_limit_per_second": 20
            }
        }
        
        # 部署配置
        self.max_concurrent_per_region = 20
        self.rate_limit_per_second = 15  # 75%的API限制
        self.batch_size = 50
        
    def create_server_in_region(self, region: str, instance_index: int, total_count: int) -> str:
        """在指定地域创建服务器"""
        try:
            region_info = self.regions[region]
            print(f"🚀 {region_info['name']} 创建服务器 {instance_index+1}/{total_count}")
            
            response = requests.post(f"{self.base_url}/api/v1/servers", 
                                   json={
                                       'region': region, 
                                       'gpu_type': 'basic'
                                   }, 
                                   timeout=120)
            
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    instance_id = data.get('instance_id')
                    print(f"✅ {region_info['name']} 服务器创建成功: {instance_id}")
                    return instance_id
                else:
                    print(f"❌ {region_info['name']} 服务器创建失败: {data.get('error')}")
                    return ""
            else:
                print(f"❌ {region_info['name']} 创建请求失败: {response.status_code}")
                return ""
        except Exception as e:
            print(f"❌ {region_info['name']} 创建异常: {e}")
            return ""
    
    def deploy_in_region(self, region: str, target_count: int) -> int:
        """在单个地域部署指定数量的服务器"""
        region_info = self.regions[region]
        print(f"\n🌍 开始在地域 {region_info['name']} 部署 {target_count} 个服务器")
        
        created_count = 0
        
        # 使用线程池并行创建
        with ThreadPoolExecutor(max_workers=self.max_concurrent_per_region) as executor:
            # 提交所有创建任务
            future_to_index = {
                executor.submit(self.create_server_in_region, region, i, target_count): i 
                for i in range(target_count)
            }
            
            # 收集结果
            for future in as_completed(future_to_index):
                index = future_to_index[future]
                try:
                    instance_id = future.result()
                    if instance_id:
                        created_count += 1
                except Exception as e:
                    print(f"❌ {region_info['name']} 创建任务 {index+1} 异常: {e}")
        
        print(f"✅ {region_info['name']} 部署完成，创建了 {created_count}/{target_count} 个服务器")
        return created_count
    
    def deploy_multi_region(self, total_count: int) -> bool:
        """多地域并行部署"""
        print(f"🌍 开始多地域并行部署 {total_count} 个服务器")
        print("=" * 60)
        
        # 计算每个地域的部署数量
        region_count = len(self.regions)
        instances_per_region = total_count // region_count
        remaining_instances = total_count % region_count
        
        print(f"📊 部署计划:")
        print(f"  总实例数: {total_count}")
        print(f"  支持地域: {region_count} 个")
        print(f"  每地域基础: {instances_per_region} 个")
        print(f"  剩余分配: {remaining_instances} 个")
        print()
        
        # 分配部署任务
        deployment_tasks = {}
        for i, (region, region_info) in enumerate(self.regions.items()):
            # 前几个地域分配剩余实例
            extra = 1 if i < remaining_instances else 0
            target_count = instances_per_region + extra
            deployment_tasks[region] = target_count
            print(f"  {region_info['name']}: {target_count} 个实例")
        
        print()
        
        # 并行部署所有地域
        total_created = 0
        start_time = time.time()
        
        with ThreadPoolExecutor(max_workers=region_count) as executor:
            # 提交所有地域的部署任务
            future_to_region = {
                executor.submit(self.deploy_in_region, region, count): region
                for region, count in deployment_tasks.items()
            }
            
            # 收集结果
            for future in as_completed(future_to_region):
                region = future_to_region[future]
                try:
                    created = future.result()
                    total_created += created
                except Exception as e:
                    print(f"❌ 地域 {region} 部署异常: {e}")
        
        end_time = time.time()
        deployment_time = end_time - start_time
        
        print(f"\n🎉 多地域部署完成！")
        print(f"📊 总计创建: {total_created}/{total_count} 个服务器")
        print(f"⏱️ 部署时间: {deployment_time:.1f} 秒")
        print(f"🚀 部署速度: {total_created/deployment_time:.1f} 个/秒")
        
        if total_created == total_count:
            print("✅ 部署成功率: 100%")
            return True
        else:
            success_rate = (total_created / total_count) * 100
            print(f"⚠️ 部署成功率: {success_rate:.1f}%")
            return False
    
    def get_system_status(self) -> Dict[str, Any]:
        """获取系统状态"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/system/status", timeout=10)
            if response.status_code == 200:
                return response.json()
            else:
                print(f"❌ 获取系统状态失败: {response.status_code}")
                return {}
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return {}
    
    def get_servers(self) -> List[Dict[str, Any]]:
        """获取所有服务器列表"""
        try:
            response = requests.get(f"{self.base_url}/api/v1/servers", timeout=10)
            if response.status_code == 200:
                data = response.json()
                if data.get('success'):
                    return data.get('data', [])
            else:
                print(f"❌ 获取服务器列表失败: {response.status_code}")
                return []
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return []
    
    def show_status(self):
        """显示系统状态"""
        print("📊 多地域部署系统状态:")
        print("=" * 50)
        
        # 系统状态
        status = self.get_system_status()
        if status and status.get('success'):
            data = status['data']
            print(f"系统运行: {data['system']['running']}")
            print(f"扩缩容运行: {data['scaling']['running']}")
            print(f"总服务器: {data['servers']['servers']['total']}")
            print(f"活跃服务器: {data['servers']['servers']['available']}")
            print(f"队列长度: {data['queue']['total_tasks_in_queue']}")
            print()
        
        # 按地域统计服务器
        servers = self.get_servers()
        if servers:
            region_stats = {}
            for server in servers:
                region = server.get('region', 'Unknown')
                if region not in region_stats:
                    region_stats[region] = 0
                region_stats[region] += 1
            
            print(f"📋 地域服务器分布 (共{len(servers)}个):")
            for region, count in region_stats.items():
                region_name = self.regions.get(region, {}).get('name', region)
                print(f"  {region_name} ({region}): {count} 个")
        else:
            print("❌ 无法获取服务器列表")

def main():
    parser = argparse.ArgumentParser(description='多地域大规模GPU服务器部署脚本')
    parser.add_argument('action', choices=['deploy', 'status'], 
                       help='操作类型')
    parser.add_argument('--count', type=int, default=300, 
                       help='要部署的服务器数量')
    parser.add_argument('--url', default='http://localhost:8080', 
                       help='HAI优化系统URL')
    
    args = parser.parse_args()
    
    manager = MultiRegionDeploymentManager(args.url)
    
    if args.action == 'deploy':
        print("🌍 多地域大规模GPU服务器部署工具")
        print("=" * 60)
        print(f"目标数量: {args.count} 个服务器")
        print(f"支持地域: {len(manager.regions)} 个")
        print(f"每地域分配: {args.count // len(manager.regions)} 个")
        print("=" * 60)
        
        # 显示地域信息
        print("📍 支持的地域:")
        for region, info in manager.regions.items():
            print(f"  {info['name']} ({region}): {info['app_id']}")
        print()
        
        # 确认部署
        confirm = input("确认开始多地域大规模部署？(y/N): ")
        if confirm.lower() != 'y':
            print("❌ 部署已取消")
            return
        
        success = manager.deploy_multi_region(args.count)
        
        if success:
            print("\n🎉 多地域大规模部署成功完成！")
        else:
            print("\n⚠️ 多地域大规模部署部分完成，请检查系统状态")
            
    elif args.action == 'status':
        manager.show_status()

if __name__ == '__main__':
    main()
